import numpy as np

#def assemble_global_stiffness(K, k_local, element, dof_per_node):
    #node1, node2 = element
   # global_dof = np.array([node1 * dof_per_node + i for i in range(dof_per_node)] +
                      #    [node2 * dof_per_node + i for i in range(dof_per_node)])
  #  for i in range(12):
     #   for j in range(12):
       #    K[global_dof[i], global_dof[j]] += k_local[i, j]


import numpy as np


def coordinate_transformation_matrix(node1_coords, node2_coords):
    """
    计算从局部坐标系到全局坐标系的转换矩阵。
    假设为3D的情况，节点1和节点2为输入的坐标。
    """
    # 计算两个节点之间的矢量
    dx = node2_coords[0] - node1_coords[0]
    dy = node2_coords[1] - node1_coords[1]
    dz = node2_coords[2] - node1_coords[2]

    L = np.sqrt(dx ** 2 + dy ** 2 + dz ** 2)  # 计算单元长度
    if L == 0:
        raise ValueError("两个节点坐标重合，无法计算转换矩阵")

    # 计算方向余弦
    l = dx / L
    m = dy / L
    n = dz / L

    # 构建旋转矩阵的正交基
    T = np.zeros((12, 12))
    T[0, 0] = T[3, 3] = T[6, 6] = T[9, 9] = l
    T[1, 1] = T[4, 4] = T[7, 7] = T[10, 10] = m
    T[2, 2] = T[5, 5] = T[8, 8] = T[11, 11] = n
    # 可以根据具体问题需要进行完整扩展

    return T


def assemble_global_stiffness(K, k_local, element, dof_per_node, nodes):
    """
    汇总曲梁单元刚度矩阵到全局刚度矩阵，考虑坐标变换
    """
    node1, node2 = element
    node1_coords = nodes[node1]
    node2_coords = nodes[node2]

    # 计算转换矩阵
    T = coordinate_transformation_matrix(node1_coords, node2_coords)

    # 变换局部刚度矩阵到全局坐标系
    k_global = T.T @ k_local @ T

    # 汇总到全局刚度矩阵
    global_dof = np.array([node1 * dof_per_node + i for i in range(dof_per_node)] +
                          [node2 * dof_per_node + i for i in range(dof_per_node)])

    for i in range(12):
        for j in range(12):
            K[global_dof[i], global_dof[j]] += k_global[i, j]
